#!/usr/bin/python3
from collections import defaultdict
from sys import argv
import time

class Node:
	def __init__(self, val):
		self.val = val
		self.in_nodes = []
		self.out_nodes = []

	def is_balanced(self):
		return len(self.out_nodes) - len(self.in_nodes) == 0

	def __hash__(self):
		return hash(self.val)

	def __eq__(self, other):
		return self.val == other.val

	def __str__(self):
		return self.val

	def __repr__(self):
		return "--"+str(self)+", "+str([str(i) for i in self.in_nodes])+", "+str([str(i) for i in self.out_nodes])+"--";

class Sbh:
	def __init__(self, lmers):
		# lmers should be list not set apparently
		lmer_length = len(next(iter(lmers)))
		assert lmers, "Cardinality of lmer set has to be at last 1"
		assert all(len(x) == lmer_length for x in lmers), "Lmers have to have same length"
		assert all(all(y in 'ATGC' for y in x) for x in lmers), "Lmers may only contain nucleotides ATGC"
		assert lmer_length >= 2, "Length of lmers has to be at least 2"
		self._path = None
		self._graph, self._start_node = self._generate_graph(lmers)

	def _generate_graph(self, lmers):
		"""
			Returns the eulerian graph and the start node of the graph
		"""
		# contains all L-1 lmers
		graph = {}
		# maps the start of each lmer (l-1 nucleotides) to all l-1 mers that it connects to
		start = defaultdict(list)
		for l in lmers:
			if l[:-1] not in graph:
				graph[l[:-1]] = Node(l[:-1])
			try:
				node = graph[l[1:]]
			except KeyError:
				node = Node(l[1:])
				graph[l[1:]] = node
			start[l[:-1]].append(node)
		for node in graph.values():
			children = start[node.val]
			for child in children:
				self._add_edge(node, child)
		start_node = self._make_eulerian(graph)
		print(graph, start_node);
		return graph.values(), start_node

	@staticmethod
	def _add_edge(start, end):
		start.out_nodes.append(end)
		end.in_nodes.append(start)

	@staticmethod
	def _make_eulerian(graph):
		"""
		We have to make sure that the two nodes whose in and out node numbers aren't equal,
		 get connected with an edge.
		Returns start node of given graph
		"""
		nodes = [n for n in graph.values() if not n.is_balanced()]
		if len(nodes) != 2:
			print(nodes);
		assert len(nodes) == 2, "Graph can only have two semi-balanced nodes"
		if len(nodes[1].in_nodes) > len(nodes[1].out_nodes):
			start = nodes[0]
			end = nodes[1]
		else:
			start = nodes[1]
			end = nodes[0]
		Sbh._add_edge(end, start)
		assert start.is_balanced() and end.is_balanced()
		return start

	def get_sequence_representation(self):
		"""
		Returns a string representation of the constructed sequence
		"""
		path = self.get_sequence()
		print(path);
		node = path[0]
		# TODO not 100% sure about this for n >= 3
		res = [node.val] + [node.val[-1] for node in path[1:]]
		return ''.join(res)[:-1]

	def get_sequence(self):
		if not self._path:
			self._path = self._compute_path(self._start_node)
		return self._path

	@staticmethod
	def _compute_path(node):
		""" Returns an eulerian cycle starting from the given node"""
		path = []
		cycle = Sbh._get_cycle(node)
		for node in cycle:
			for n in node.out_nodes:
				path.extend(Sbh._compute_path(n))
			path.append(node)
		return path

	@staticmethod
	def _get_cycle(node):
		"""
		Finds a cycle in the graph starting from the given node and returns the complete cycle.
		"""
		path = []
		while True:
			try:
				print("Adding "+repr(node)+" to path.");
				path.append(node)
				node = node.out_nodes.pop(0)
			except IndexError:
				# we finished the cycle
				break
		return path

def parse_sequences(lines):
	# *should* be a set according to algorithm description, but that leads to incorrect results.
	res = []
	sequence = []
	for line in lines:
		if line.startswith('>') or line.startswith(';'):
			if sequence: res.append(''.join(sequence))
			sequence = []
		else:
			sequence.append(line.upper())
	if sequence: res.append(''.join(sequence))
	return res

def read_file(filename):
	with open(filename, 'r', encoding='UTF-8') as file:
		return parse_sequences(file.read().splitlines())

def write_file(filename, sequence):
	# fasta file format says that line length shouldn't be longer than
	# 120 characters, but 80 is generally preferred
	MAX_LINE_LENGTH = 80
	with open(filename, 'w', encoding='UTF-8') as file:
		file.write('> sequence created by SBH approach\n')
		for i in range(0, len(sequence), MAX_LINE_LENGTH):
			file.write(sequence[i:i+MAX_LINE_LENGTH])
			file.write('\n')

if __name__ == '__main__':
	start_time = time.time()
	segments = read_file(argv[1])
	sbh = Sbh(segments)
	seq = sbh.get_sequence_representation()
	end_time = time.time()

	print(seq)
	original = ''.join(open("./data/SBH_orig_seq"+argv[1][len("./data/SBH_reads"):]).read().replace("\r","").split("\n")[1:])
	print(len(original))
	print(len(seq))
	write_file(argv[1][:-len(".fasta")]+"-assembled.fasta", seq)
	open(argv[1][:-len(".fasta")]+"-timing",'w',encoding='UTF-8').write(str(end_time - start_time))
	open(argv[1][:-len(".fasta")]+"-score",'w',encoding='UTF-8').write(str(sum([(1 if seq[i] != (original[i] if i < len(original) else ' ') else 0) for i in range(len(seq))])))
	
